from django.shortcuts import redirect

from mainapp.models import User


def my_login_required(view_func):
    def _wrapped_view_func(request, *args, **kwargs):
        username = request.session.get('username')
        if username and User.objects.filter(username=username).exists():
            return view_func(request, *args, **kwargs)
        else:
            return redirect('login')

    return _wrapped_view_func


def my_login_admin_required(view_func):
    def _wrapped_view_func(request, *args, **kwargs):
        role = request.session.get('role')
        if role and role == 'admin':
            return view_func(request, *args, **kwargs)
        else:
            return redirect('login')

    return _wrapped_view_func
